from copy import deepcopy
from typing import Union
from enum import Enum

import numpy as np

import coptpy as cp
from coptpy import COPT


# L前缀表示该变量对应的拉格朗日乘子
class terminal_args:
    class args_status(Enum):
        SKIP = 0
        INUSE = 1
        CONSTANT = 2

    def __init__(self, T: int, require=("P", "Q", "V"), constant=()):
        self.require = require
        self.constant = constant
        # 端口有功，指向设备为正
        if "P" in require:
            self.P = np.zeros((T, 1), dtype=np.float64)
            if "P" not in constant:
                self.LP = np.zeros((T, 1), dtype=np.float64)
        # 端口无功，指向设备为正
        if "Q" in require:
            self.Q = np.zeros((T, 1), dtype=np.float64)
            if "Q" not in constant:
                self.LQ = np.zeros((T, 1), dtype=np.float64)
        # 端口电压平方
        if "V" in require:
            self.V2 = np.zeros((T, 1), dtype=np.float64)
            if "V" not in constant:
                self.LV2 = np.zeros((T, 1), dtype=np.float64)

    def is_total_constant(self) -> bool:
        status_list = [self.var_status("P"), self.var_status("Q"), self.var_status("V")]
        b = True
        for s in status_list:
            b = b and (s == self.args_status.CONSTANT or s == self.args_status.SKIP)
        return b

    def var_status(self, v_str: str) -> args_status:
        if v_str in self.require:
            if v_str in self.constant:
                return self.args_status.CONSTANT
            else:
                return self.args_status.INUSE
        else:
            return self.args_status.SKIP

    def todict(self) -> dict:
        d = {}
        if self.var_status("P") == self.args_status.INUSE:
            d.update(
                {"P": self.P.reshape(-1).tolist(), "LP": self.LP.reshape(-1).tolist()}
            )
        elif self.var_status("P") == self.args_status.CONSTANT:
            d.update({"P": self.P.reshape(-1).tolist(), "LP": None})
        if "Q" in self.require:
            if "Q" in self.constant:
                d.update({"Q": self.Q.reshape(-1).tolist(), "LQ": None})
            else:
                d.update(
                    {
                        "Q": self.Q.reshape(-1).tolist(),
                        "LQ": self.LQ.reshape(-1).tolist(),
                    }
                )
        if "V" in self.require:
            if "V" in self.constant:
                d.update({"V": self.V2.reshape(-1).tolist(), "LV": None})
            else:
                d.update(
                    {
                        "V2": self.V2.reshape(-1).tolist(),
                        "LV2": self.LV2.reshape(-1).tolist(),
                    }
                )
        return d


def node_update(
    adj_terminals: list[terminal_args], rou: float, T: int
) -> list[terminal_args]:
    """
    根据k步的拉格朗日乘子、k+1步端口状态yn_k+1，返回更新后的k+1步的节点/端口状态zn_k+1
    adj_terminals 包含k步的拉格朗日乘子、k+1步的状态变量的端口状态
    rou 惩罚因子
    """

    alpha = np.zeros((T, 1))
    beta = np.zeros((T, 1))
    ita = np.zeros((T, 1))
    terminal_num = {"P": 0, "Q": 0, "V": 0}
    z_adj_terminals = deepcopy(adj_terminals)
    # 式（38）更新功率
    for t in adj_terminals:
        if t.var_status("P") == t.args_status.INUSE:
            alpha += t.LP + rou * t.P
            terminal_num["P"] += 1
        elif t.var_status("P") == t.args_status.CONSTANT:
            alpha += rou * t.P
        if t.var_status("Q") == t.args_status.INUSE:
            beta += t.LQ + rou * t.Q
            terminal_num["Q"] += 1
        elif t.var_status("Q") == t.args_status.CONSTANT:
            beta += rou * t.Q
    alpha /= terminal_num["P"]
    beta /= terminal_num["Q"]

    for z in z_adj_terminals:
        if z.var_status("P") == z.args_status.INUSE:
            z.P = z.P + (z.LP - alpha) / rou
        if z.var_status("Q") == z.args_status.INUSE:
            z.Q = z.Q + (z.LQ - beta) / rou
        if z.var_status("V") == z.args_status.INUSE:
            z.V2 = ita / rou

    # 式（38）更新电压
    has_constant_v=False #标志是否有恒定电压的消息
    constant_v2=None
    for t in adj_terminals:
        if t.var_status("V") == t.args_status.CONSTANT:
            # 先禁止传入两个恒定电压消息，就算一样也不行
            if has_constant_v:
                raise RuntimeError("禁止传入两个恒定电压，就算一样电压也不行")
            else:
                has_constant_v=True
                constant_v2 = t.V2

    if not has_constant_v:
        # 常规更新流程
        for t in adj_terminals:
            if t.var_status("V") == t.args_status.INUSE:
                ita += t.LV2 + rou * t.V2
                terminal_num["V"] += 1
        ita /= terminal_num["V"]
        for z in z_adj_terminals:
            if z.var_status("V") == z.args_status.INUSE:
                z.V2 = ita / rou
    else:
        # 直接赋值所有电压
        for z in z_adj_terminals:
            if z.var_status("V") == z.args_status.INUSE:
                z.V2 = constant_v2
    return z_adj_terminals


def acline_update(
    R: float,
    X: float,
    Vmax: float,
    Vmin: float,
    Imax: float,
    terminal1: terminal_args,
    terminal2: terminal_args,
    rou: float,
    T: int,
) -> dict:
    """
    交流线路，无决策变量
    只有状态变量：两端有功，电压幅值平方，电流幅值平方
    """
    # Create COPT environment
    env = cp.Envr()

    # Create COPT model
    model = env.createModel("device_acline")

    # Add variables
    Pi = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Pi",
    )
    Pj = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Pj",
    )
    Qi = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Qi",
    )
    Qj = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Qj",
    )
    if Imax is None:
        l2 = model.addMVar(
            (T, 1),
            lb=0.0,
            ub=COPT.INFINITY,
            vtype=COPT.CONTINUOUS,
            nameprefix="I2",
        )
    else:
        l2 = model.addMVar(
            (T, 1), lb=0.0, ub=Imax**2, vtype=COPT.CONTINUOUS, nameprefix="I2"
        )
    V2i = model.addMVar(
        (T, 1),
        lb=Vmin**2,
        ub=Vmax**2,
        vtype=COPT.CONTINUOUS,
        nameprefix="V2i",
    )
    V2j = model.addMVar(
        (T, 1),
        lb=Vmin**2,
        ub=Vmax**2,
        vtype=COPT.CONTINUOUS,
        nameprefix="V2j",
    )
    half_l2 = model.addMVar(
        (T, 1),
        lb=0.0,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="halfI",
    )
    """
    cnorm2 = model.addVar(lb=0.0, ub=COPT.INFINITY, name="cnorm2")
    auxPi = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="auxPi",
    )
    auxPj = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="auxPj",
    )
    auxQi = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Qi",
    )
    auxQj = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Qj",
    )
    auxV2i = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="V2i",
    )
    auxV2j = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="V2j",
    )
    """
    # Add linear constraints
    model.addConstrs(Pi + Pj == l2 * R, nameprefix="CP")
    model.addConstrs(Qi + Qj == l2 * X, nameprefix="CQ")
    model.addConstrs(
        V2j == V2i - 2 * (Pi * R + Qi * X) + l2 * (R**2 + X**2),
        nameprefix="CV",
    )
    # Add auxiliary varibles constraints
    model.addConstrs(half_l2 == l2 / 2.0, nameprefix="ChalfI")
    """
    model.addConstrs(auxPi == Pi - terminal1.P, nameprefix="CauxPi")
    model.addConstrs(auxQi == Qi - terminal1.Q, nameprefix="CauxQi")
    model.addConstrs(auxV2i == V2i - terminal1.V2, nameprefix="CauxV2i")
    model.addConstrs(auxPj == Pj - terminal2.P, nameprefix="CauxPj")
    model.addConstrs(auxQj == Qj - terminal2.Q, nameprefix="CauxQj")
    model.addConstrs(auxV2j == V2j - terminal2.V2, nameprefix="CauxV2j")
    """
    # Add 锥约束
    for k in range(0, T):
        model.addCone(
            [half_l2[k].item(), V2i[k].item(), Pi[k].item(), Qi[k].item()],
            COPT.CONE_RQUAD,
        )
    # 惩罚二次项就写成旋转二阶锥
    """
    CONE_cvars = [cnorm2, cnorm2]
    CONE_cvars.extend(auxPi.tolist())
    CONE_cvars.extend(auxQi.tolist())
    CONE_cvars.extend(auxV2i.tolist())
    CONE_cvars.extend(auxPj.tolist())
    CONE_cvars.extend(auxQj.tolist())
    CONE_cvars.extend(auxV2j.tolist())
    model.addCone(CONE_cvars, COPT.CONE_RQUAD)
    """

    # Set linear objective
    obj = (
        terminal1.LP.transpose() @ Pi
        + terminal1.LQ.transpose() @ Qi
        + terminal1.LV2.transpose() @ V2i
        + terminal2.LP.transpose() @ Pj
        + terminal2.LQ.transpose() @ Qj
        + terminal2.LV2.transpose() @ V2j
    )
    """
    obj = obj + (rou / 2.0) * (2.0 * cnorm2)
    obj = obj + rou / 2.0 * (
        (Pi - terminal1.P) @ (Pi - terminal1.P)
        + (Pj - terminal2.P).transpose() @ (Pj - terminal2.P)
        + (Qi - terminal1.Q).transpose() @ (Qi - terminal1.Q)
        + (Qj - terminal2.Q).transpose() @ (Qj - terminal2.Q)
        + (V2i - terminal1.V2).transpose() @ (V2i - terminal1.V2)
        + (V2j - terminal2.V2).transpose() @ (V2j - terminal2.V2)
    )
    """
    obj = obj + rou / 2.0 * (
        (
            Pi.transpose() @ Pi
            - 2 * terminal1.P.transpose() @ Pi
            + terminal1.P.transpose() @ terminal1.P
        )
        + (
            Qi.transpose() @ Qi
            - 2 * terminal1.Q.transpose() @ Qi
            + terminal1.Q.transpose() @ terminal1.Q
        )
        + (
            V2i.transpose() @ V2i
            - 2 * terminal1.V2.transpose() @ V2i
            + terminal1.V2.transpose() @ terminal1.V2
        )
        + (
            Pj.transpose() @ Pj
            - 2 * terminal2.P.transpose() @ Pj
            + terminal2.P.transpose() @ terminal2.P
        )
        + (
            Qj.transpose() @ Qj
            - 2 * terminal2.Q.transpose() @ Qj
            + terminal2.Q.transpose() @ terminal2.Q
        )
        + (
            V2j.transpose() @ V2j
            - 2 * terminal2.V2.transpose() @ V2j
            + terminal2.V2.transpose() @ terminal2.V2
        )
    )
    model.setObjective(obj, COPT.MINIMIZE)

    # Set parameters
    # model.setParam(COPT.Param.TimeLimit, 7200)
    model.setParam(COPT.Param.Logging, 0)

    # Solve the problem
    model.solve()

    # gather solution and update y
    # 返回得到的k+1步交流线路端口状态和第k+1步的拉格朗日乘子
    y = {
        "I2": l2.x,
        "Pi": Pi.x,
        "Qi": Qi.x,
        "V2i": V2i.x,
        "Pj": Pj.x,
        "Qj": Qj.x,
        "V2j": V2j.x,
        "LPi": terminal1.LP + rou * (Pi.x - terminal1.P),
        "LQi": terminal1.LQ + rou * (Qi.x - terminal1.Q),
        "LV2i": terminal1.LV2 + rou * (V2i.x - terminal1.V2),
        "LPj": terminal2.LP + rou * (Pj.x - terminal2.P),
        "LQj": terminal2.LQ + rou * (Qj.x - terminal2.Q),
        "LV2j": terminal2.LV2 + rou * (V2j.x - terminal2.V2),
        "gap": V2i.x * l2.x - (Pi.x * Pi.x + Qi.x * Qi.x),
    }
    if model.status != COPT.OPTIMAL:
        # raise RuntimeError("交流线路更新未能找到最优解")
        y["optimal_converged"]=False
    else:
        y["optimal_converged"]=True
    return y


def upper_grid_update(
    Ct: np.ndarray,
    Pmax: float,
    Pmin: float,
    terminal: terminal_args,
    rou: float,
    T: int,
) -> dict:
    # Create COPT environment
    env = cp.Envr()
    # Create COPT model
    model = env.createModel("device_upper_grid")
    # Add variables
    Psub = model.addMVar(
        (T, 1),
        lb=Pmin,
        ub=Pmax,
        vtype=COPT.CONTINUOUS,
        nameprefix="Psub",
    )
    Pi = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Pi",
    )
    # Add auxiliary varibles constraints
    model.addConstrs(Pi == -1.0 * Psub, nameprefix="CauxPi")

    # Set objective
    obj = Ct.transpose() @ Psub + terminal.LP.transpose() @ Pi
    obj = obj + rou / 2.0 * (
        Pi.transpose() @ Pi
        - 2 * terminal.P.transpose() @ Pi
        + terminal.P.transpose() @ terminal.P
    )

    model.setObjective(obj, COPT.MINIMIZE)
    # Set parameters
    model.setParam(COPT.Param.Logging, 0)
    # Solve the problem
    model.solve()

    # gather solution and update y and x
    if model.status != COPT.OPTIMAL:
        raise RuntimeError("上级电网更新未能找到最优解")
    # 无功最优解为等于端口无功
    y = {
        "Psub": Psub.x,
        "Qsub": -1.0 * terminal.Q,
        "Pi": Pi.x,
        "Qi": terminal.Q,
        "Cupper": Ct.transpose() @ Psub.x,
        "LPi": terminal.LP + rou * (Pi.x - terminal.P),
        "LQi": np.zeros((T, 1)),
    }
    return y


def storage_update(
    PEssMax: float,
    eff_ch: float,
    eff_dis: float,
    capacity: float,
    max_changing_N: int,
    cost_operating: float,
    depreciation_per_change: float,
    interval_seconds: float,
    terminal: terminal_args,
    rou: float,
    T: int,
) -> dict:
    """
    PEssMax 最大充放电功率,kWh或MWh
    eff_ch 充电效率, 0到1
    eff_dis 放电效率, 0到1
    capacity 容量 kWh或MWh
    interval_seconds两个时间点之间间隔的秒数
    T必须大于2
    """
    # Create COPT environment
    env = cp.Envr()
    # Create COPT model
    model = env.createModel("device_storage")
    # Add variables
    Pch = model.addVars(
        T,
        lb=0,
        ub=PEssMax,
        vtype=COPT.CONTINUOUS,
        nameprefix="Pch",
    )
    Pdis = model.addVars(
        T,
        lb=0,
        ub=PEssMax,
        vtype=COPT.CONTINUOUS,
        nameprefix="Pch",
    )
    state_ch = model.addVars(
        T,
        lb=0,
        ub=1,
        vtype=COPT.BINARY,
        nameprefix="state_ch",
    )
    change_var = model.addVars(
        T - 1,
        lb=0,
        ub=1,
        vtype=COPT.BINARY,
        nameprefix="change",
    )
    Ee = model.addVars(
        T,
        lb=0.1 * capacity,
        ub=0.9 * capacity,
        vtype=COPT.CONTINUOUS,
        nameprefix="Ee",
    )

    Pi = model.addVars(
        T,
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Pch",
    )

    # Add linear constraints
    model.addConstrs(
        [Pch[t] <= state_ch[t] * PEssMax for t in range(0, T)],
        nameprefix="ChargeCons",
    )
    model.addConstrs(
        [Pdis[t] <= (1 - state_ch[t]) * PEssMax for t in range(0, T)],
        nameprefix="DischargeCons",
    )
    model.addConstrs(
        [
            Ee[t]
            == Ee[t - 1]
            + eff_ch * Pch[t - 1] * (interval_seconds / 3600)
            - eff_dis * Pdis[t - 1] / eff_dis * (interval_seconds / 3600)
            for t in range(1, T)
        ],
        nameprefix="EeCons",
    )
    model.addConstr(
        Ee[0] == Ee[T - 1],
        name="EeSameC",
    )
    model.addConstrs(
        [change_var[t - 1] >= (state_ch[t] - state_ch[t - 1]) for t in range(1, T)],
        nameprefix="change_b1_cons",
    )
    model.addConstrs(
        [change_var[t - 1] >= (state_ch[t - 1] - state_ch[t]) for t in range(1, T)],
        nameprefix="change_b2_cons",
    )
    model.addConstr(change_var.sum() <= max_changing_N, name="changing_limit")
    model.addConstrs(
        [Pi[t] == Pch[t] - Pdis[t] for t in range(0, T)], nameprefix="power_sum_cons"
    )

    # Set linear objective
    obj = (Pch.sum() + Pdis.sum()) * cost_operating * (interval_seconds / 3600)

    if depreciation_per_change is not None:
        obj = obj + depreciation_per_change * change_var.sum()
    else:
        depreciation_per_change = 0.0

    for t in range(0, T):
        obj = obj + terminal.LP[t] * Pi[t]

    # 二次项
    for t in range(0, T):
        obj = obj + rou / 2 * (Pi[t] - terminal.P[t]) * (Pi[t] - terminal.P[t])

    model.setObjective(obj, COPT.MINIMIZE)

    # Set parameters
    model.setParam(COPT.Param.Logging, 0)

    # Solve the problem
    model.solve()

    # gather solution and update y
    # 返回得到的k+1步交流线路端口状态和第k步的拉格朗日乘子
    if model.status != COPT.OPTIMAL:
        raise RuntimeError("储能更新未能找到最优解")
    else:
        # 我的结论是这种数组变量没有矩阵变量好用
        y = {
            "Pch": np.array([Pch[k].x for k in range(0, T)], ndmin=2).transpose(),
            "Pdis": np.array([Pdis[k].x for k in range(0, T)], ndmin=2).transpose(),
            "Pi": np.array([Pi[k].x for k in range(0, T)], ndmin=2).transpose(),
            "Qi": terminal.Q,
            "state_ch": np.array(
                [state_ch[k].x for k in range(0, T)], ndmin=2
            ).transpose(),
            "Ee": np.array([Ee[k].x for k in range(0, T)], ndmin=2).transpose(),
            "Cost": (
                np.sum(np.array([Pch[k].x for k in range(0, T)]))
                + np.sum(np.array([Pdis[k].x for k in range(0, T)]))
            )
            * cost_operating
            * (interval_seconds / 3600)
            + depreciation_per_change
            * np.sum(np.array([change_var[k].x for k in range(0, T-1)])),
        }
        y.update(
            {"LPi": terminal.LP + rou * (y["Pi"] - terminal.P), "LQi": np.zeros((T, 1))}
        )
    return y


def load_preprocess(
    data: Union[int, float, np.ndarray], expected_slot: int
) -> np.ndarray:
    if isinstance(data, int):
        return np.ones((expected_slot, 1), dtype=np.int32) * data
    elif isinstance(data, float):
        return np.ones((expected_slot, 1), dtype=np.float64) * data
    elif isinstance(data, np.ndarray):
        if np.size(data) >= expected_slot:
            return np.reshape(data[0:expected_slot], (expected_slot, 1))
        else:
            y = np.zeros((expected_slot, 1), dtype=np.float64)
            round = np.size(data)
            count = 0
            while expected_slot - count >= round:
                y[count : count + round] = data
                count += round
            y[count:expected_slot] = data[0 : expected_slot - count]
            return y
    else:
        raise TypeError("Unsupported data type")


def load_update(
    Pt: np.ndarray,
    Qt: np.ndarray,
    terminal: terminal_args,
    rou: float,
    T: int,
) -> terminal_args:
    """
    当前仅考虑恒定功率负荷
    最优解就是保持功率
    """
    y = terminal_args(T, ("P", "Q"), ("P", "Q"))
    y.P = load_preprocess(Pt, T)
    y.LP = terminal.LP + rou * (y.P - terminal.P)
    y.Q = load_preprocess(Qt, T)
    y.LQ = terminal.LQ + rou * (y.Q - terminal.Q)
    return y
